import argparse

import os
import inspect
import sys
import numpy as np
from importlib import import_module
import paddle.fluid as fluid

from hhutil.io import fmt_path, is_hidden, rename, rm

import helm.pretrained
from helm.static.models import calculate_params


def parse_args():
    def str2bool(v):
        return v.lower() in ("true", "t", "1")

    parser = argparse.ArgumentParser()
    parser.add_argument("-m", "--model", type=str)
    parser.add_argument("-d", "--download", type=str)
    parser.add_argument("-p", "--ppcls", type=str)
    parser.add_argument("-t", "--transpile", default=False, action='store_true')

    return parser.parse_args()


if __name__ == '__main__':
    args = parse_args()

    ppcls_dir = fmt_path(args.ppcls)
    model_name = args.model
    download_dir = fmt_path(args.download)

    output_dir = download_dir / (model_name + "_output")
    output_dir.mkdir(exist_ok=True)
    export_dir = download_dir / (model_name + "_pretrained")
    model_filename = output_dir / "model"
    params_filename = output_dir / "params"

    python_exec = sys.executable
    def run_ppcls(cmd):
        os.system(f"cd {ppcls_dir} && PYTHONPATH=.:$PYTHONPATH python {cmd}")

    if not export_dir.exists():
        run_ppcls(f"./tools/download.py -a {model_name} -p {download_dir} -d True")
    if not params_filename.exists():
        run_ppcls(f"./tools/export_model.py -m {model_name} -p {export_dir} -o {output_dir}")

    place = fluid.CPUPlace()
    exe = fluid.Executor(place)
    [program, feed_names, fetch_list] = fluid.io.load_inference_model(
        str(output_dir), exe,
        model_filename=str(model_filename),
        params_filename=str(params_filename))

    rm(model_filename)
    rm(params_filename)

    g = fluid.global_scope()
    state_dict = {}
    for f in export_dir.iterdir():
        if not is_hidden(f):
            name = f.name
            v = g.find_var(name)
            if v is not None:
                state_dict[name] = np.array(v.get_tensor())

    fluid.enable_imperative()
    sd = state_dict
    for k, v in state_dict.items():
        sd[k] = fluid.dygraph.to_variable(v, name=k, zero_copy=True)

    if args.transpile and hasattr(helm.pretrained, model_name):
        model = getattr(helm.pretrained, model_name)()
        transpile = getattr(import_module(model.__module__), "transpile")
        if len(inspect.signature(transpile).parameters) == 2:
            new_sd = transpile(sd, model)
        else:
            new_sd = transpile(sd)
        n = 0
        for k, v in new_sd.items():
            n += np.prod(v.shape)
        assert n == calculate_params(model)[0]
        model.set_dict(new_sd)
        sd = model.state_dict()
        fluid.dygraph.save_dygraph(sd, str(output_dir / (model_name + "_transpile")))
        rm(export_dir)
        rm(output_dir / (model_name + ".pdparams"))
    else:
        fluid.dygraph.save_dygraph(sd, str(output_dir / model_name))
        rename(output_dir / (model_name + ".pdopt"), model_name + ".pdparams", stem=False)
    fluid.disable_imperative()